
import onnx ,onnxruntime as ort
import numpy as np

class onnx_model:
    def __init__(self,onnx_model_path) -> None: 
        # torch.onnx.export(traced, *adapter.flattened_inputs, onnx_model_path, verbose=True, input_names=['obs'], output_names=['mu','sigma', 'value'])
        onnx_model = onnx.load(onnx_model_path)

        # Check that the model is well formed
        onnx.checker.check_model(onnx_model)
        self.ort_model = ort.InferenceSession(onnx_model_path)

    def infference(self,obs):
        outputs = self.ort_model.run(None,{'obs' : obs})

        mu = outputs[0]
        sigma = np.exp(outputs[1])
        # action = np.random.normal(mu, sigma)
        
        # xjz
        action=mu

        
        return action
